import random
import pygame
from pygame.locals import *
import os

# 定义一些颜色变量，背景色
BLACK = pygame.Color(0, 0, 0)
RED = pygame.Color(255, 0, 0)
GREEN = pygame.Color(0, 255, 0)
WHITE = pygame.Color(255, 255, 255)
HEAD_COLOR = pygame.Color(0, 119, 255)
BORDER_COLOR = pygame.Color(50, 50, 50)  # 边界颜色

# 全局配置变量
HEIGHT = 32  # 高度
WIDTH = 32  # 宽度
CELL_SIZE = 20  # 每个单元格的像素大小
MAX_STEPS = 100000  # 最大步数
SNAKE_SPEED = 1500  # 蛇的速度，帧率

FIELD_SIZE = HEIGHT * WIDTH  # 字段大小等于长乘以宽
HEAD = 0  # 定义蛇头的位置 位于snake数组的第一个位置

# 用数字代表不同的对象，运动时矩阵上每个格子会处理成到达食物的路径长度
UNDEFINED = (HEIGHT + 1) * (WIDTH + 1)
SNAKE = 2 * UNDEFINED
FOOD = 0  # 食物

# 由于snake是一维数组，所以对应元素直接加上以下值就表示向四个方向移动
LEFT = -1
RIGHT = 1
UP = -WIDTH  # 一维数组，所以需要整个宽度都加上才能表示上下移动
DOWN = WIDTH

# 错误码
ERR = -1111

# 用一维数组来表示二维的东西
board = [0] * FIELD_SIZE
snake = [0] * (FIELD_SIZE + 1)
snake[HEAD] = 1 * WIDTH + 1  # 初始化蛇头在(1,1)的地方
snake_size = 1  # 定义蛇的长度为1

# 与上面变量对应的临时变量，蛇试探性地移动时使用，既在蛇真正的开始运动之前，先模拟一次
tmpboard = [0] * FIELD_SIZE
tmpsnake = [0] * (FIELD_SIZE + 1)
tmpsnake[HEAD] = 1 * WIDTH + 1
tmpsnake_size = 1

food = 4 * WIDTH + 7  # 食物位置初始在(4, 7)
best_move = ERR  # 运动方向，最终确定的方向

# 运动方向数组，游戏分数(蛇长)
mov = [LEFT, RIGHT, UP, DOWN]  # 蛇能够运动的四个方向

score = 1  # 分数，蛇触碰到食物，score就加一
high_score_file = "high_score.txt"
high_score = 0
step_counter = 0  # 步数计数器

# 初始化pygame
pygame.init()
window_width = WIDTH * CELL_SIZE + 200  # 增加宽度以显示分数
window_height = HEIGHT * CELL_SIZE
fpsClock = pygame.time.Clock()
playSurface = pygame.display.set_mode((window_width, window_height))  # 设置窗口大小
pygame.display.set_caption('贪吃蛇')

# 设置字体
font = pygame.font.SysFont("Arial", 18)

def load_high_score():
    """加载历史最高分"""
    global high_score
    if os.path.exists(high_score_file):
        with open(high_score_file, "r") as file:
            high_score = int(file.read().strip())
    else:
        high_score = 0

def save_high_score(new_high_score):
    """保存新的历史最高分"""
    with open(high_score_file, "w") as file:
        file.write(str(new_high_score))

def display_scores():
    """显示当前分数和历史最高分"""
    current_score_text = font.render(f"Score: {score}", True, WHITE)
    high_score_text = font.render(f"High Score: {high_score}", True, WHITE)
    playSurface.blit(current_score_text, (WIDTH * CELL_SIZE + 20, 10))  # 调整分数显示位置
    playSurface.blit(high_score_text, (WIDTH * CELL_SIZE + 20, 30))

def draw_wall():
    """绘制墙壁"""
    for i in range(WIDTH):
        draw_cell(i, RED)
        draw_cell(i + (HEIGHT - 1) * WIDTH, RED)
    for i in range(HEIGHT):
        draw_cell(i * WIDTH, RED)
        draw_cell((i + 1) * WIDTH - 1, RED)

def draw_cell(position, color):
    """绘制单个方块，并添加边界"""
    x, y = CELL_SIZE * (position % WIDTH), CELL_SIZE * (position // WIDTH)
    pygame.draw.rect(playSurface, color, Rect(x, y, CELL_SIZE, CELL_SIZE))
    pygame.draw.rect(playSurface, BORDER_COLOR, Rect(x, y, CELL_SIZE, CELL_SIZE), 1)  # 添加边界

def draw_snake_and_food():
    """绘制蛇和食物"""
    playSurface.fill(BLACK)  # 清空整个屏幕
    draw_wall()  # 绘制墙壁
    draw_cell(food, RED)  # 绘制食物
    for i in range(snake_size):  # 绘制蛇
        color = HEAD_COLOR if i == 0 else WHITE
        draw_cell(snake[i], color)
    display_scores()  # 显示当前分数和历史最高分
    pygame.display.flip()

def is_cell_free(idx, psize, psnake):
    """检查一个cell有没有被蛇身覆盖，没有覆盖则为free，返回true"""
    return not (idx in psnake[:psize])

def is_move_possible(idx, move):
    """检查某个位置idx是否可向move方向运动"""
    if move == LEFT:
        return idx % WIDTH > 1
    elif move == RIGHT:
        return idx % WIDTH < (WIDTH - 2)
    elif move == UP:
        return idx > (2 * WIDTH - 1)
    elif move == DOWN:
        return idx < (FIELD_SIZE - 2 * WIDTH)
    return False

def board_reset(psnake, psize, pboard):
    """重置board"""
    for i in range(FIELD_SIZE):
        if i == food:
            pboard[i] = FOOD
        elif is_cell_free(i, psize, psnake):
            pboard[i] = UNDEFINED
        else:
            pboard[i] = SNAKE

def board_BFS(pfood, psnake, pboard):
    """广度优先搜索遍历整个board，计算出board中每个非SNAKE元素到达食物的路径长度"""
    queue = [pfood]
    inqueue = [0] * FIELD_SIZE
    found = False
    while queue:
        idx = queue.pop(0)
        if inqueue[idx] == 1:
            continue
        inqueue[idx] = 1
        for move in mov:
            if is_move_possible(idx, move):
                if idx + move == psnake[HEAD]:
                    found = True
                if pboard[idx + move] < SNAKE:
                    if pboard[idx + move] > pboard[idx] + 1:
                        pboard[idx + move] = pboard[idx] + 1
                    if inqueue[idx + move] == 0:
                        queue.append(idx + move)
    return found

def choose_shortest_safe_move(psnake, pboard):
    """从蛇头开始，根据board中元素值，从蛇头周围4个领域点中选择最短路径"""
    min_distance = SNAKE
    best_move = ERR
    for move in mov:
        if is_move_possible(psnake[HEAD], move) and pboard[psnake[HEAD] + move] < min_distance:
            min_distance = pboard[psnake[HEAD] + move]
            best_move = move
    return best_move

def choose_longest_safe_move(psnake, pboard):
    """从蛇头开始，根据board中元素值，从蛇头周围4个领域点中选择最远路径"""
    max_distance = -1
    best_move = ERR
    for move in mov:
        if is_move_possible(psnake[HEAD], move) and pboard[psnake[HEAD] + move] < UNDEFINED and pboard[psnake[HEAD] + move] > max_distance:
            max_distance = pboard[psnake[HEAD] + move]
            best_move = move
    return best_move

def is_tail_inside():
    """检查是否可以追着蛇尾运动，即蛇头和蛇尾间是有路径的"""
    tmpboard[tmpsnake[tmpsnake_size - 1]] = 0
    tmpboard[food] = SNAKE
    result = board_BFS(tmpsnake[tmpsnake_size - 1], tmpsnake, tmpboard)
    for move in mov:
        if is_move_possible(tmpsnake[HEAD], move) and tmpsnake[HEAD] + move == tmpsnake[tmpsnake_size - 1] and tmpsnake_size > 3:
            result = False
    return result

def follow_tail():
    """让蛇头朝着蛇尾运行一步"""
    tmpsnake_size = snake_size
    tmpsnake[:] = snake[:]
    board_reset(tmpsnake, tmpsnake_size, tmpboard)
    tmpboard[tmpsnake[tmpsnake_size - 1]] = FOOD
    tmpboard[food] = SNAKE
    board_BFS(tmpsnake[tmpsnake_size - 1], tmpsnake, tmpboard)
    tmpboard[tmpsnake[tmpsnake_size - 1]] = SNAKE
    return choose_longest_safe_move(tmpsnake, tmpboard)

def any_possible_move():
    """在各种方案都不行时，随便找一个可行的方向来走(1步)"""
    board_reset(snake, snake_size, board)
    board_BFS(food, snake, board)
    min_distance = SNAKE
    best_move = ERR
    for move in mov:
        if is_move_possible(snake[HEAD], move) and board[snake[HEAD] + move] < min_distance:
            min_distance = board[snake[HEAD] + move]
            best_move = move
    return best_move

def shift_array(arr, size):
    """转换数组函数"""
    for i in range(size, 0, -1):
        arr[i] = arr[i - 1]

def new_food():
    """随机生成新的食物"""
    global food, snake_size
    cell_free = False
    while not cell_free:
        w = random.randint(1, WIDTH - 2)
        h = random.randint(1, HEIGHT - 2)
        food = WIDTH * h + w
        cell_free = is_cell_free(food, snake_size, snake)
    draw_cell(food, RED)

def make_move(pbest_move):
    """真正的蛇在这个函数中，朝pbest_move走1步"""
    global key, snake, board, snake_size, score, step_counter
    shift_array(snake, snake_size)
    snake[HEAD] += pbest_move
    p = snake[HEAD]

    draw_snake_and_food()
    step_counter += 1

    if snake[HEAD] == food:
        board[snake[HEAD]] = SNAKE
        snake_size += 1
        score += 1
        if snake_size < FIELD_SIZE:
            new_food()
    else:
        board[snake[HEAD]] = SNAKE
        board[snake[snake_size]] = UNDEFINED

def virtual_shortest_move():
    """虚拟运行吃到食物后，得到虚拟下蛇在board的位置"""
    global snake, board, snake_size, tmpsnake, tmpboard, tmpsnake_size, food
    tmpsnake_size = snake_size
    tmpsnake[:] = snake[:]
    tmpboard[:] = board[:]
    board_reset(tmpsnake, tmpsnake_size, tmpboard)

    food_eated = False
    while not food_eated:
        board_BFS(food, tmpsnake, tmpboard)
        move = choose_shortest_safe_move(tmpsnake, tmpboard)
        shift_array(tmpsnake, tmpsnake_size)
        tmpsnake[HEAD] += move
        if tmpsnake[HEAD] == food:
            tmpsnake_size += 1
            board_reset(tmpsnake, tmpsnake_size, tmpboard)
            tmpboard[food] = SNAKE
            food_eated = True
        else:
            tmpboard[tmpsnake[HEAD]] = SNAKE
            tmpboard[tmpsnake[tmpsnake_size]] = UNDEFINED

def find_safe_way():
    """如果蛇与食物间有路径，则调用本函数"""
    global snake, board
    safe_move = ERR
    virtual_shortest_move()
    if is_tail_inside():
        return choose_shortest_safe_move(snake, board)
    safe_move = follow_tail()
    return safe_move

def game_loop():
    """游戏主循环"""
    global high_score
    load_high_score()
    while True:
        board_reset(snake, snake_size, board)
        if board_BFS(food, snake, board):
            best_move = find_safe_way()
        else:
            best_move = follow_tail()
        if best_move == ERR:
            best_move = any_possible_move()
        if best_move != ERR:
            make_move(best_move)
        else:
            break
        if step_counter > MAX_STEPS:
            break
        fpsClock.tick(SNAKE_SPEED)  # 控制游戏速度

    # 检查是否需要更新最高分
    if score > high_score:
        high_score = score
        save_high_score(high_score)

    print("Game Over! Final Score:", score)

if __name__ == "__main__":
    new_food()  # 初始化食物位置
    game_loop()  # 开始游戏循环